import cv2
import subprocess
import time

# RTSP流地址和RTMP服务器地址
RTSP_URL = "rtsp://telchina:123qweASD@10.10.30.39:554//Streaming/Channels/1"
RTMP_URL = "rtmp://10.10.50.127:10935/live/stream_5"

def stream_video():
    while True:  # 持续尝试推送视频流
        try:
            cap = cv2.VideoCapture(RTSP_URL)

            if not cap.isOpened():
                print("无法打开RTSP流，等待后重试...")
                time.sleep(5)
                continue

            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            fps = cap.get(cv2.CAP_PROP_FPS)

            command = [
                'ffmpeg',
                '-y',
                '-f', 'rawvideo',
                '-pix_fmt', 'bgr24',
                '-s', f'{width}x{height}',
                '-r', str(int(fps)),  # 使用RTSP流的实际帧率
                '-i', '-',  # 从stdin读取视频流
                '-c:v', 'libx264',
                '-pix_fmt', 'yuv420p',
                '-preset', 'ultrafast',
                '-tune', 'zerolatency',
                '-b:v', '1000k',
                '-maxrate', '1000k',
                '-bufsize', '2000k',
                '-c:a', 'aac',
                '-b:a', '128k',  # 指定音频比特率
                '-ac', '2',  # 指定音频通道数
                '-f', 'flv',
                RTMP_URL
            ]

            process = subprocess.Popen(command, stdin=subprocess.PIPE, bufsize=10**8)

            while True:
                ret, frame = cap.read()
                
                if not ret:
                    print("读取帧失败，结束流处理")
                    break

                # 绘制识别框
                cv2.rectangle(frame, (50, 50), (200, 200), (0, 255, 0), 2)

                # 显示当前帧（可选）
                cv2.imshow("Video Stream", frame)

                # 将帧写入ffmpeg进程的stdin
                process.stdin.write(frame.tobytes())

                # 按 'q' 键退出
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    break

            process.stdin.close()
            process.kill()

        except Exception as e:
            print(f"发生异常: {e}")
        
        finally:
            cap.release()
            cv2.destroyAllWindows()
            print("尝试重新推送视频流...")
            time.sleep(5)

if __name__ == "__main__":
    stream_video()